
import numpy as np
import cv2
import open3d as o3d
from Camera import Camera




class PointGenerator:
    """
    生成点云数据
    """
    def __init__(self,disparity,img_color):
        if isinstance(img_color, np.ndarray):
            self.depth_map = disparity
            self.img_color = cv2.cvtColor(img_color, cv2.COLOR_BGR2RGB)
        else:
            # 从文件路径加载视差图和彩色图
            self.depth_map = cv2.imread(disparity, cv2.IMREAD_UNCHANGED)
            self.img_color = cv2.cvtColor(cv2.imread(img_color), cv2.COLOR_BGR2RGB)
        size = (640, 480)
        file_path = "./相机信息/out.xls"
        self.camera = Camera(file_path, size)
        self.height, self.width = self.depth_map.shape
        #生成点云数据
        self.points = self.generate3DOriginal(True)[0]
        self.colors = self.generate3DOriginal(True)[1]

    def generate3DOriginal(self, tag=False):
        """
        根据tag的值，选择返回点云数据还是点云可视化
        """
        colors, points = self.ThreeDOriginal()

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        pcd.colors = o3d.utility.Vector3dVector(colors)

        if tag:
            return (points,colors)

        o3d.visualization.draw_geometries([pcd])

    def ThreeDOriginal(self):
        """遍历视差图，根据视差计算每个像素的三维坐标"""
        points = []
        colors = []
        for v in range(self.height):
            for u in range(self.width):
                D = self.depth_map[v, u]
                #过滤无效视差图
                if D == 0 or D < 40:
                    continue
                #计算Z深度值
                Z = (self.camera.b * self.camera.f_x) / D
                if Z > 0:
                    #计算X、Y的值
                    X = self.camera.b * (u - self.camera.cx) / D
                    Y = self.camera.b * (v - self.camera.cy) / D

                    color = self.img_color[v, u] / 255.0

                    points.append([X, Y, Z])
                    colors.append(color)
        points = np.array(points)
        colors = np.array(colors)
        return colors, points
